package com.example.demo.service;

/**
 * @author qinbo
 * @version 1.0
 * @description: TODO
 * @date 2025/2/11 7:22
 */
import com.example.demo.properties.OllamaProperties;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;
import java.util.Map;

@Service
public class OllamaChatService {

    private final WebClient webClient;
    private final OllamaProperties properties;

    public OllamaChatService(WebClient.Builder webClientBuilder, OllamaProperties properties) {
        this.webClient = webClientBuilder.baseUrl(properties.getBaseUrl()).build();
        this.properties = properties;
    }

    public Mono<String> generate(String prompt) {
        if (!properties.getChat().isEnabled()) {
            return Mono.error(new IllegalStateException("Ollama chat is disabled"));
        }

        Map<String, Object> requestBody = Map.of(
                "model", properties.getChat().getModel(),
                "prompt", prompt,
                "stream", false,
                "options", properties.getChat().getOptions()
        );

        return webClient.post()
                .uri("/api/generate")
                .bodyValue(requestBody)
                .retrieve()
                .onStatus(status -> status.isError(), response ->
                        response.bodyToMono(String.class)
                                .flatMap(error -> Mono.error(new RuntimeException("Ollama API Error: " + error)))
                )
                .bodyToMono(Map.class)
                .map(response -> (String) response.get("response"));
    }
}
